﻿
// 2022/7/4: 首个版本
// 2022/7/5: 修正writeToDB3File
// 2022/11/11: 增加writeToDB3File和DB3FileReader::initialize返回值
// 2023/10/9: 添加英文注释
// 2025/5/16: 修正注释

#ifndef ROS2_SUPPORT_H
#define ROS2_SUPPORT_H

#include <memory>
#include "rcutils/types.h"
#include "rosidl_generator_c/message_type_support_struct.h"
#include "rosidl_generator_cpp/traits.hpp"
#include "ros2_support_base.h"

// Need to import libraries / 需要添加库依赖: rcutils, rmw_fastrtps_cpp

extern "C"
{
    int32_t rmw_serialize(const void *ros_message, const rosidl_message_type_support_t *type_support, rcutils_uint8_array_t *serialized_message);
    int32_t rmw_deserialize(const rcutils_uint8_array_t *serialized_message, const rosidl_message_type_support_t *type_support, void *ros_message);
}

namespace spadas
{
    /// \~English @brief Write ROS2 data object to the db3 file
    /// \~Chinese @brief 将ROS2数据对象写入db3文件
    /// \~English @param caller Cross-module function caller interface
    /// \~Chinese @param caller 跨模块函数调用接口
    /// \~English @param rawFolder Whether to write to the db3 file in the raw data folder (input), otherwise it is the generation data folder (output)
    /// \~Chinese @param rawFolder 是否写至原始数据文件夹(input)内的db3文件，否则为generation数据文件夹(output)
    /// \~English @param fileName File name (do not need to include the suffix ".db3")
    /// \~Chinese @param fileName 文件名(不需要包括后缀.db3)
    /// \~English @param topic Topic
    /// \~Chinese @param topic Topic
    /// \~English @param time The time offset, in seconds
    /// \~Chinese @param time 相对时间戳，单位秒
    /// \~English @param obj ROS2 data object
    /// \~Chinese @param obj ROS2数据对象
    /// \~English @param typeSupport ROS2 data type support
    /// \~Chinese @param typeSupport ROS2数据类型支持
    /// \~English @returns Whether the data is successfully written
    /// \~Chinese @returns 是否成功写入
    template<typename Type>
    Bool writeToDB3File(Interface<ICrossCaller> caller, Bool rawFolder, String fileName, String topic, Double time, Type& obj, const rosidl_message_type_support_t *typeSupport)
    {
        auto serialized = std::make_shared<rcutils_uint8_array_t>();
        auto defaultAllocator = rcutils_get_default_allocator();
        if (rcutils_uint8_array_init(serialized.get(), 0, &defaultAllocator)) {}
        int32_t res = rmw_serialize(&obj, typeSupport, serialized.get());
        Bool ok = FALSE;
        if (res == 0)
        {
            Ros2FileIOWriteMessageObject ioObject;
            ioObject.setWriteToGeneration(!rawFolder);
            ioObject.setFileName(fileName);
            ioObject.setTopic(topic);
            ioObject.setTime(time);
            ioObject.setDataType(rosidl_generator_traits::data_type<Type>());
            ioObject.setBinaryData(Binary(serialized->buffer, serialized->buffer_length));
            ok = caller->callNativeFunction("Ros2FileIO", "WriteMessage", ioObject.as<BaseObject>());
        }
        if (rcutils_uint8_array_fini(serialized.get())) {}
        return ok;
    }

    /// \~English @brief One ROS2 data frame
    /// \~Chinese @brief 一帧ROS2数据
    template<typename Type>
    struct DB3Object
    {
        /// \~English @brief The time offset, in seconds
        /// \~Chinese @brief 相对时间戳，单位秒
        Double time;

        /// \~English @brief ROS2 data object
        /// \~Chinese @brief ROS2数据对象
        Type obj;
    };

    /// \~English @brief Read ROS2 data from the db3 file
    /// \~Chinese @brief 从db3文件读取ROS2数据
    template<typename Type>
    class DB3FileReader
    {
    public:
        /// \~English @brief Initialization
        /// \~Chinese @brief 初始化
        /// \~English @param caller Cross-module function caller interface
        /// \~Chinese @param caller 跨模块函数调用接口
        /// \~English @param rawFolder Whether to read data from db3 file in raw data folder (input), otherwise it is generation data folder (output)
        /// \~Chinese @param rawFolder 是否从原始数据文件夹(input)内的db3文件读取数据，否则为generation数据文件夹(output)
        /// \~English @param fileName File name (do not need to include the suffix ".db3")
        /// \~Chinese @param fileName 文件名(不需要包括后缀.db3)
        /// \~English @param topic Topic
        /// \~Chinese @param topic Topic
        /// \~English @param streamSize Size of the buffer queue
        /// \~Chinese @param streamSize 缓存队列的大小
        /// \~English @param typeSupport ROS2 data type support
        /// \~Chinese @param typeSupport ROS2数据类型支持
        /// \~English @returns Whether the initialization is successful
        /// \~Chinese @returns 是否成功初始化
        Bool initialize(Interface<ICrossCaller> caller, Bool rawFolder, String fileName, String topic, UInt streamSize, const rosidl_message_type_support_t *typeSupport)
        {
            stream = Stream<Ros2FileIOReadMessage>(math::max(1u, streamSize));

            Ros2FileIORegisterReadStreamObject ioObject;
            ioObject.setReadFromGeneration(!rawFolder);
            ioObject.setFileName(fileName);
            ioObject.setTopic(topic);
            ioObject.setReadStream(stream);

            if (!caller->callNativeFunction("Ros2FileIO", "RegisterReadStream", ioObject.as<BaseObject>()))
            {
                stream.terminate();
                return FALSE;
            }

            this->typeSupport = typeSupport;
            return TRUE;
        }

        /// \~English @brief Retrieves all new data from the buffer queue
        /// \~Chinese @brief 从缓存队列中取出所有新数据
        Array<DB3Object<Type> > dequeue()
        {
            auto newMessages = stream.dequeue();
            if (newMessages.isEmpty()) return Array<DB3Object<Type> >();

            Array<DB3Object<Type> > output = Array<DB3Object<Type> >::createUninitialized(newMessages.size());
            for (UInt i = 0; i < newMessages.size(); i++)
            {
                Ros2FileIOReadMessage msg = newMessages[i];
                DB3Object<Type> newObject;
                newObject.time = msg.time;

                auto serialized = std::make_shared<rcutils_uint8_array_t>();
                auto defaultAllocator = rcutils_get_default_allocator();
                if (rcutils_uint8_array_init(serialized.get(), msg.binaryData.size(), &defaultAllocator)) {}
                utility::memoryCopy(msg.binaryData.data(), serialized->buffer, msg.binaryData.size());
                serialized->buffer_length = msg.binaryData.size();
                rmw_deserialize(serialized.get(), typeSupport, &newObject.obj);

                output.initialize(i, newObject);
            }
            return output;
        }

        /// \~English @brief Whether it has been read to the end
        /// \~Chinese @brief 是否已读至末尾
        Bool endOfFile()
        {
            return stream.isTerminated();
        }

    private:
        Stream<Ros2FileIOReadMessage> stream;
        const rosidl_message_type_support_t *typeSupport;
    };
}

#endif